Lambda関数でTensorFlow Liteをインポートし、推論処理を実行する

山本紘暉

2021.06.02

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

はじめに

現在、カフェのシステムでは、機械学習を用いて、カメラを用いて動画を撮影し、商品の前にいる人物の骨格や手を検出することで、どのユーザがどの商品を取り出したかを判定しています。

今までは、骨格検出モデルを用いてエッジデバイスで動画を推論処理（撮影した画像から映っている人物の骨格の座標を検出する処理）を実行する、という構成で処理をしていました。今後、エッジ側のデバイスの費用を下げたり、骨格検出以外の処理を増やすことを考えているため、エッジデバイスからクラウドに動画を送信し、クラウド側で様々な処理を実行する、という構成を検討しています。

前回までの記事で、エッジデバイスでの動画処理（エンコード・送信）と、クラウド側の処理（動画の取り出し）について記載しました。

撮影した動画をリアルタイムにエンコードする方法【GStreamer】

【Kinesis Video Streams】Pythonで動画ファイルを送信する

Lambda関数でKinesis Video Streamsから動画を取得する

課題

AWSで機械学習のモデルを使った推論処理を実行するには、SageMaker Endpointを利用する方法も考えられますが、Endpointは立ち上げるのに5分程度かかるため、判定をリアルタイムに行い、ユーザに表示するには長すぎ、常時立ち上げておく必要があります。

しかし、性能が低めのインスタンスタイプ（ml.c5.large）でも、常時立ち上げると￥7200/mon程度かかり、結構高い料金がかかります。カフェでは、商品を取ったときのみ処理をすれば良く、常時処理する必要がないため、無駄が多く、節約できる余地がありそうです。

目的

今回は、クラウド側のメインの処理となる、機械学習のモデルを使った推論処理を、料金を節約しながら行うことを目的とします。そのために、Lambda関数で推論処理を動かす方法、特に、Lambda関数でTensorFlow Liteを動かすことで、ある程度速い処理速度で動かす方法を記載します。

TFLiteを動かす

Lambda関数でTensorFlow Lite（以下、TFLite）を動かすには、tflite_runtimeをインストールしたLambdaレイヤを作成し、インポートする必要があります。

レイヤを作成する

結論

このリポジトリで公開しているようなライブラリを作成します。layertflite5C7B319D-e3a04c3e-3f55-4014-a044-619fbb991d5f.zipを、Lambdaのレイヤにアップロードすれば作成可能です（レイヤのランタイムはpython3.7を選択してください）。このzipファイルは、libraryファイル内をzipで圧縮したものです。

https://github.com/cm-yamamoto-hiroki/aws_lambda_layer_tflite

作成方法

先程のライブラリの作成方法について説明します。

今回の手順ではDocker環境が必要です。手元に環境がある場合には、そちらをお使いいただいて問題ありません。環境がない場合、EC2でインスタンスを作成すると簡単です。今回は、EC2インスタンスを新たに立ち上げた場合について書きます。

EC2の設定は、以下のようにしました

Amazon マシンイメージ (AMI)：Amazon Linux 2 AMI (HVM), SSD Volume Type, 64 ビット (x86)
インスタンスタイプの選択：t2.xlarge（t2.microだと処理時間が長いです。料金がかかりますが、少し性能の高いインスタンスタイプがおすすめです。）

キーペアは既存のものを利用しても良いですし、新しく作成しても構いません。インスタンスに接続したら、以下のコマンドを実行しました。今回は、こちらのリポジトリを利用させていだたきました。（途中、自分が動作確認したコミットでブランチを作成しています）

S3に出力する場合は、インスタンスのIAMロールにS3へのアクセス権限のあるポリシーをアタッチしておく必要があります。また、スクリプト中の「python3.7」は、EC2のpython3のバージョンに合わせて変更してください。

sudo yum install -y git docker
sudo usermod -a -G docker ec2-user

cd
git clone https://github.com/tpaul1611/python_tflite_for_amazonlinux
cd python_tflite_for_amazonlinux
git checkout -b temp 415ec188df3862514a0cd9a088a6d639201ba78b

sudo service docker start
sudo docker build -t tflite_amazonlinux .
sudo docker run -d --name=tflite_amazonlinux tflite_amazonlinux
sudo docker cp tflite_amazonlinux:/usr/local/lib64/python3.7/site-packages .
sudo docker stop tflite_amazonlinux

mkdir -p ./python/lib/python3.7
mv site-packages ./python/lib/python3.7
zip -r layer.zip python/

# S3に出力する場合：
# aws s3 cp layer.zip s3://mybucketname/

# 接続元PCに取り出す場合：接続を切り、scpなどで取り出す

取り出したライブラリ（zipファイル）を、Lambdaのレイヤに追加します。ランタイムはEC2のpython3のバージョンを選択してください。

推論処理を実行する

Lambda関数を作成します。ランタイムは上のものと合わせてください。モデルの推論処理は（一般的に）処理量が多いので、メモリサイズを大きめ（1024MB～など）にし、タイムアウトも長め（30s～など）に設定する必要があります。

コードとしては、以下のように実装しました。コンストラクタのfilepath_model_checkpointに、TFLite用に変換したファイルへのパスを引数として与えて生成し、inferに画像を入力することで、推論処理の結果を得られます。モデルファイルは、予めLambda関数のパッケージに含めておくか、S3バケットに保存しておいて、使用前に/tmp以下などにコピーしておく必要があります。入力する画像は、サイズをモデルの入力サイズに合わせ、かつ、モデルの入力の前提（明度の正規化など）に合わせておく必要があります。

from typing import List

import time
import numpy as np
import tflite_runtime.interpreter as tflite

import logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

def load_model_checkpoint(filepath_model_checkpoint):
    interpreter = tflite.Interpreter(filepath_model_checkpoint)
    # メモリ確保。これはモデル読み込み直後に必須
    interpreter.allocate_tensors()

    return interpreter

class TfLiteModel:
    def __init__(self, filepath_model_checkpoint):
        self._filepath_model_checkpoint = filepath_model_checkpoint
        self._interpreter = load_model_checkpoint(filepath_model_checkpoint)

    def infer(self, img_processed: np.ndarray) -> List[np.ndarray]:
        # 学習モデルの入力層・出力層のプロパティをGet.
        input_details = self._interpreter.get_input_details()

        # 入力層のテンソルデータ構成の取得
        input_shape = input_details[0]['shape']
        input_dtype = input_details[0]['dtype']

        # 推論処理を実行
        input_data = img_processed.astype(input_dtype)
        self._interpreter.set_tensor(input_details[0]['index'], input_data)
        self._interpreter.invoke()

        # 推論結果は、output_detailsのindexに保存されている
        output_details = self._interpreter.get_output_details()
        output_data = [
            self._interpreter.get_tensor(output_details[i]['index'])[0]
            for i in reversed(range(len(output_details))  # torchの出力形式に合わせる
        )]
        for i in range(len(output_data)):
            data, detail = output_data[i], output_details[len(output_details)-1 - i]
            logger.info([detail["name"], data.shape])

        return output_data

    @property
    def model_input_shape(self):
        shape = self._interpreter.get_input_details()[0]["shape"]  # ex) [1, 3, 236, 416] (=[n_image, n_depth, n_height, n_width])
        width, height = shape[3], shape[2]
        return width, height

注意点としては、インポートは「import tflite_runtime.interpreter」まで行う必要があります。「import tflite_runtime」とすると、「tflite_runtime.interpreter」としたところで、エラーが発生しました。

参考の値ですが、骨格検出のモデルを動かした際、メモリサイズを10240MBに設定して実行し、416*236の画像を入力した結果、「self._interpreter.invoke()」の処理時間は3.0s程度でした（SageMaker Endpointを利用した場合、0.1秒程度で終わりました）。Lambda関数は並列して実行できるため、多くの画像があっても、劇的に処理時間が伸びることはなさそうです。

例えば、商品を取るユーザの動画が3秒（=画像30枚）である場合、処理にかかる時間料金は￥1.5程度です。数百回程度処理しても、SageMaker Endpointを立ち上げるよりも、大きく料金を抑えることができます。（実際に利用する際には、予めメモリサイズを変えながら何回か実行して処理時間を計測し、料金を最適化できるサイズを選択した方が良さそうです）

まとめ

クラウドでの推論処理にかかる料金を抑えるために、Lambda関数で推論処理を実行する方法を調べました。TensorFlow LiteをインストールしたLambdaレイヤを作成し、Lambda関数からインポートして推論処理を実行することで、ある程度の処理速度を維持したまま、料金を抑えることができることがわかりました。